import React, { useRef, useState, useEffect } from 'react';
import { ProTable, ActionType, ProColumns } from '@ant-design/pro-components';
import { message } from 'antd';
import api from '../../services/api';
import { Agent, TokenUsage, User } from '../../types';

const TokenUsageList: React.FC = () => {
  const actionRef = useRef<ActionType>(null);
  const [users, setUsers] = useState<Record<number, User>>({});
  const [agents, setAgents] = useState<Record<number, Agent>>({});

  useEffect(() => {
    fetchUsersAndAgents();
  }, []);

  const fetchUsersAndAgents = async () => {
    try {
      const usersRes = await api.get('/users/', { params: { skip: 0, limit: 1000 } });

      const usersMap: Record<number, User> = {};
      usersRes.data.items.forEach((user: User) => {
        usersMap[user.id] = user;
      });

      // Fetch agents for display
      try {
        const agentRes = await api.get('/agents/', { params: { skip: 0, limit: 1000 } });
        const agentMap: Record<number, Agent> = {};
        agentRes.data.items.forEach((agent: Agent) => {
          agentMap[agent.id] = agent;
        });
        setAgents(agentMap);
      } catch (e) {
        // ignore; optional data
      }

      setUsers(usersMap);
    } catch (error: any) {
      message.error('Failed to fetch reference data');
    }
  };

  const fetchTokenUsage = async (params: any) => {
    try {
      const response = await api.get('/token-usage/', {
        params: {
          skip: (params.current - 1) * params.pageSize,
          limit: params.pageSize,
          user_id: params.user_id,
          agent_id: params.agent_id,
        },
      });
      return {
        data: response.data.items,
        success: true,
        total: response.data.total,
      };
    } catch (error: any) {
      message.error('Failed to fetch token usage data');
      return {
        data: [],
        success: false,
        total: 0,
      };
    }
  };

  const columns: ProColumns<TokenUsage>[] = [
    {
      title: 'ID',
      dataIndex: 'id',
      width: 80,
      search: false,
    },
    {
      title: 'User',
      dataIndex: 'user_id',
      render: (_, record) => users[record.user_id]?.username || `User ${record.user_id}`,
      valueType: 'select',
      valueEnum: Object.fromEntries(
        Object.entries(users).map(([id, user]) => [id, { text: user.username }])
      ),
    },
    {
      title: 'Agent',
      dataIndex: 'agent_id',
      render: (_, record) => agents[record.agent_id]?.name || `Agent ${record.agent_id}`,
      valueType: 'select',
      valueEnum: Object.fromEntries(
        Object.entries(agents).map(([id, agent]) => [id, { text: agent.name }])
      ),
    },
    {
      title: 'Tokens Used',
      dataIndex: 'tokens_used',
      search: false,
      sorter: true,
    },
    {
      title: 'Cost ($)',
      dataIndex: 'cost',
      search: false,
      render: (_, record) => `$${record.cost.toFixed(4)}`,
      sorter: true,
    },
    {
      title: 'Created At',
      dataIndex: 'created_at',
      valueType: 'dateTime',
      search: false,
      sorter: true,
    },
  ];

  return (
    <ProTable<TokenUsage>
      columns={columns}
      actionRef={actionRef}
      request={fetchTokenUsage}
      rowKey="id"
      search={{
        labelWidth: 'auto',
      }}
      pagination={{
        defaultPageSize: 10,
        showSizeChanger: true,
      }}
      dateFormatter="string"
      headerTitle="Token Usage Analytics"
    />
  );
};

export default TokenUsageList;
